"""
多线程管理器 - 提供数据库系统的多线程支持
"""

import os
import sys
import time
import queue
import threading
from typing import Dict, List, Optional, Callable, Any, Tuple
from concurrent.futures import ThreadPoolExecutor, Future
from dataclasses import dataclass
from enum import Enum
import uuid

# 添加当前目录到Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
if current_dir not in sys.path:
    sys.path.insert(0, current_dir)

try:
    from statistics import Logger
except ImportError:
    from .statistics import Logger


class TaskType(Enum):
    """任务类型枚举"""
    QUERY = "query"
    INSERT = "insert"
    UPDATE = "update"
    DELETE = "delete"
    IO_READ = "io_read"
    IO_WRITE = "io_write"
    CACHE_CLEANUP = "cache_cleanup"
    STATS_UPDATE = "stats_update"


class TaskPriority(Enum):
    """任务优先级"""
    LOW = 1
    NORMAL = 2
    HIGH = 3
    CRITICAL = 4


@dataclass
class Task:
    """任务对象"""
    task_id: str
    task_type: TaskType
    priority: TaskPriority
    func: Callable
    args: tuple
    kwargs: dict
    session_id: str = None
    created_time: float = None
    
    def __post_init__(self):
        if self.created_time is None:
            self.created_time = time.time()
    
    def __lt__(self, other):
        """支持优先级排序"""
        return self.priority.value > other.priority.value


class ThreadSafeCounter:
    """线程安全计数器"""
    
    def __init__(self, initial_value: int = 0):
        self._value = initial_value
        self._lock = threading.RLock()
    
    def increment(self, amount: int = 1) -> int:
        with self._lock:
            self._value += amount
            return self._value
    
    def decrement(self, amount: int = 1) -> int:
        with self._lock:
            self._value -= amount
            return self._value
    
    def get(self) -> int:
        with self._lock:
            return self._value
    
    def set(self, value: int):
        with self._lock:
            self._value = value


class ThreadSession:
    """线程池会话管理（区别于数据库会话）"""
    
    def __init__(self, session_id: str, client_info: str = ""):
        self.session_id = session_id
        self.client_info = client_info
        self.created_time = time.time()
        self.last_activity = time.time()
        self.active_tasks = set()
        self.lock = threading.RLock()
        self.is_active = True
    
    def update_activity(self):
        """更新最后活动时间"""
        with self.lock:
            self.last_activity = time.time()
    
    def add_task(self, task_id: str):
        """添加任务到会话"""
        with self.lock:
            self.active_tasks.add(task_id)
            self.update_activity()
    
    def remove_task(self, task_id: str):
        """从会话中移除任务"""
        with self.lock:
            self.active_tasks.discard(task_id)
            self.update_activity()
    
    def close(self):
        """关闭会话"""
        with self.lock:
            self.is_active = False


class DatabaseThreadPool:
    """数据库专用线程池"""
    
    def __init__(self, 
                 query_workers: int = 4,
                 io_workers: int = 2,
                 maintenance_workers: int = 1,
                 max_queue_size: int = 1000):
        """
        初始化数据库线程池
        
        Args:
            query_workers: 查询处理线程数
            io_workers: I/O处理线程数
            maintenance_workers: 维护任务线程数
            max_queue_size: 最大队列大小
        """
        self.query_workers = query_workers
        self.io_workers = io_workers
        self.maintenance_workers = maintenance_workers
        self.max_queue_size = max_queue_size
        
        # 不同类型的线程池
        self.query_executor = ThreadPoolExecutor(
            max_workers=query_workers,
            thread_name_prefix="DB-Query"
        )
        
        self.io_executor = ThreadPoolExecutor(
            max_workers=io_workers,
            thread_name_prefix="DB-IO"
        )
        
        self.maintenance_executor = ThreadPoolExecutor(
            max_workers=maintenance_workers,
            thread_name_prefix="DB-Maintenance"
        )
        
        # 任务队列（按优先级排序）
        self.task_queue = queue.PriorityQueue(maxsize=max_queue_size)
        
        # 线程池会话管理
        self.sessions: Dict[str, ThreadSession] = {}
        self.sessions_lock = threading.RLock()
        
        # 统计信息
        self.stats = {
            'tasks_submitted': ThreadSafeCounter(),
            'tasks_completed': ThreadSafeCounter(),
            'tasks_failed': ThreadSafeCounter(),
            'active_sessions': ThreadSafeCounter(),
            'total_sessions': ThreadSafeCounter()
        }
        
        # 运行状态
        self.is_running = True
        self.logger = Logger()
        
        # 启动任务分发器
        self.dispatcher_thread = threading.Thread(
            target=self._task_dispatcher,
            name="DB-TaskDispatcher",
            daemon=True
        )
        self.dispatcher_thread.start()
        
        # 启动会话清理器
        self.session_cleaner = threading.Thread(
            target=self._session_cleaner,
            name="DB-SessionCleaner",
            daemon=True
        )
        self.session_cleaner.start()
        
        self.logger.info("THREAD_POOL", 0, 
                        f"Database thread pool initialized: "
                        f"query={query_workers}, io={io_workers}, maintenance={maintenance_workers}")
    
    def create_session(self, client_info: str = "") -> str:
        """创建新的连接会话"""
        session_id = str(uuid.uuid4())
        
        with self.sessions_lock:
            session = ThreadSession(session_id, client_info)
            self.sessions[session_id] = session
            
            self.stats['active_sessions'].increment()
            self.stats['total_sessions'].increment()
        
        self.logger.info("SESSION", 0, f"Created session {session_id[:8]}...")
        return session_id
    
    def close_session(self, session_id: str) -> bool:
        """关闭连接会话"""
        with self.sessions_lock:
            if session_id in self.sessions:
                session = self.sessions[session_id]
                session.close()
                
                # 等待会话中的任务完成
                with session.lock:
                    if session.active_tasks:
                        self.logger.info("SESSION", 0, 
                                        f"Waiting for {len(session.active_tasks)} tasks to complete...")
                
                del self.sessions[session_id]
                self.stats['active_sessions'].decrement()
                
                self.logger.info("SESSION", 0, f"Closed session {session_id[:8]}...")
                return True
        
        return False
    
    def submit_task(self, 
                   task_type: TaskType,
                   func: Callable,
                   args: tuple = (),
                   kwargs: dict = None,
                   priority: TaskPriority = TaskPriority.NORMAL,
                   session_id: str = None) -> Optional[str]:
        """
        提交任务到线程池
        
        Args:
            task_type: 任务类型
            func: 要执行的函数
            args: 函数参数
            kwargs: 函数关键字参数
            priority: 任务优先级
            session_id: 会话ID
        
        Returns:
            str: 任务ID，如果提交失败返回None
        """
        if not self.is_running:
            return None
        
        if kwargs is None:
            kwargs = {}
        
        task_id = str(uuid.uuid4())
        task = Task(
            task_id=task_id,
            task_type=task_type,
            priority=priority,
            func=func,
            args=args,
            kwargs=kwargs,
            session_id=session_id
        )
        
        try:
            # 检查队列是否已满
            if self.task_queue.full():
                self.logger.warn("TASK_QUEUE", 0, "Task queue is full, rejecting task")
                return None
            
            # 添加任务到会话
            if session_id:
                with self.sessions_lock:
                    if session_id in self.sessions:
                        self.sessions[session_id].add_task(task_id)
            
            # 提交任务到队列
            self.task_queue.put(task, timeout=1.0)
            self.stats['tasks_submitted'].increment()
            
            self.logger.info("TASK", 0, f"Submitted {task_type.value} task {task_id[:8]}...")
            return task_id
            
        except queue.Full:
            self.logger.error("TASK_QUEUE", 0, "Failed to submit task: queue full")
            return None
    
    def _task_dispatcher(self):
        """任务分发器（在独立线程中运行）"""
        self.logger.info("DISPATCHER", 0, "Task dispatcher started")
        
        while self.is_running:
            try:
                # 从队列获取任务
                task = self.task_queue.get(timeout=1.0)
                
                # 选择合适的线程池
                if task.task_type in [TaskType.IO_READ, TaskType.IO_WRITE]:
                    executor = self.io_executor
                elif task.task_type in [TaskType.CACHE_CLEANUP, TaskType.STATS_UPDATE]:
                    executor = self.maintenance_executor
                else:
                    executor = self.query_executor
                
                # 提交任务执行
                future = executor.submit(self._execute_task, task)
                
                # 标记任务队列中的任务已处理
                self.task_queue.task_done()
                
            except queue.Empty:
                continue
            except Exception as e:
                self.logger.error("DISPATCHER", 0, f"Task dispatcher error: {e}")
        
        self.logger.info("DISPATCHER", 0, "Task dispatcher stopped")
    
    def _execute_task(self, task: Task):
        """执行单个任务"""
        start_time = time.time()
        
        try:
            self.logger.info("TASK_EXEC", 0, 
                           f"Executing {task.task_type.value} task {task.task_id[:8]}...")
            
            # 执行任务
            result = task.func(*task.args, **task.kwargs)
            
            # 记录成功
            execution_time = time.time() - start_time
            self.stats['tasks_completed'].increment()
            
            self.logger.info("TASK_EXEC", 0, 
                           f"Task {task.task_id[:8]}... completed in {execution_time:.3f}s")
            
            return result
            
        except Exception as e:
            # 记录失败
            execution_time = time.time() - start_time
            self.stats['tasks_failed'].increment()
            
            self.logger.error("TASK_EXEC", 0, 
                            f"Task {task.task_id[:8]}... failed after {execution_time:.3f}s: {e}")
            raise
            
        finally:
            # 从会话中移除任务
            if task.session_id:
                with self.sessions_lock:
                    if task.session_id in self.sessions:
                        self.sessions[task.session_id].remove_task(task.task_id)
    
    def _session_cleaner(self):
        """会话清理器（清理超时会话）"""
        self.logger.info("SESSION_CLEANER", 0, "Session cleaner started")
        
        while self.is_running:
            try:
                current_time = time.time()
                timeout_sessions = []
                
                with self.sessions_lock:
                    for session_id, session in self.sessions.items():
                        # 检查会话是否超时（30分钟无活动）
                        if current_time - session.last_activity > 1800:  # 30分钟
                            timeout_sessions.append(session_id)
                
                # 清理超时会话
                for session_id in timeout_sessions:
                    self.close_session(session_id)
                    self.logger.info("SESSION_CLEANER", 0, 
                                   f"Cleaned up timeout session {session_id[:8]}...")
                
                # 每5分钟检查一次
                time.sleep(300)
                
            except Exception as e:
                self.logger.error("SESSION_CLEANER", 0, f"Session cleaner error: {e}")
                time.sleep(60)  # 出错后等待1分钟
        
        self.logger.info("SESSION_CLEANER", 0, "Session cleaner stopped")
    
    def get_stats(self) -> Dict[str, Any]:
        """获取线程池统计信息"""
        return {
            'thread_pools': {
                'query_workers': self.query_workers,
                'io_workers': self.io_workers,
                'maintenance_workers': self.maintenance_workers
            },
            'queue_status': {
                'queue_size': self.task_queue.qsize(),
                'max_queue_size': self.max_queue_size
            },
            'task_stats': {
                'submitted': self.stats['tasks_submitted'].get(),
                'completed': self.stats['tasks_completed'].get(),
                'failed': self.stats['tasks_failed'].get(),
                'success_rate': (
                    self.stats['tasks_completed'].get() / 
                    max(1, self.stats['tasks_submitted'].get())
                )
            },
            'session_stats': {
                'active_sessions': self.stats['active_sessions'].get(),
                'total_sessions': self.stats['total_sessions'].get()
            }
        }
    
    def wait_completion(self, timeout: float = None):
        """等待所有任务完成"""
        self.task_queue.join()
    
    def shutdown(self, wait: bool = True):
        """关闭线程池"""
        self.logger.info("SHUTDOWN", 0, "Shutting down database thread pool...")
        
        self.is_running = False
        
        if wait:
            # 等待队列中的任务完成
            self.wait_completion()
        
        # 关闭所有线程池
        self.query_executor.shutdown(wait=wait)
        self.io_executor.shutdown(wait=wait)
        self.maintenance_executor.shutdown(wait=wait)
        
        # 关闭所有会话
        with self.sessions_lock:
            for session_id in list(self.sessions.keys()):
                self.close_session(session_id)
        
        self.logger.info("SHUTDOWN", 0, "Database thread pool shutdown complete")


class ThreadSafeStorageEngine:
    """线程安全的存储引擎包装器"""
    
    def __init__(self, storage_engine, thread_pool: DatabaseThreadPool):
        self.storage_engine = storage_engine
        self.thread_pool = thread_pool
        self.logger = Logger()
        
        # 为存储引擎操作添加线程安全锁
        self._operation_lock = threading.RLock()
    
    def execute_async(self, operation: str, *args, session_id: str = None, **kwargs) -> Optional[str]:
        """异步执行存储引擎操作"""
        
        # 将操作包装为线程安全的函数
        def safe_operation():
            with self._operation_lock:
                method = getattr(self.storage_engine, operation)
                return method(*args, **kwargs)
        
        # 根据操作类型确定任务类型和优先级
        task_type = TaskType.QUERY
        priority = TaskPriority.NORMAL
        
        if operation.startswith('insert'):
            task_type = TaskType.INSERT
        elif operation.startswith('select') or operation.startswith('get'):
            task_type = TaskType.QUERY
        elif operation.startswith('update'):
            task_type = TaskType.UPDATE
        elif operation.startswith('delete') or operation.startswith('drop'):
            task_type = TaskType.DELETE
            priority = TaskPriority.HIGH
        
        return self.thread_pool.submit_task(
            task_type=task_type,
            func=safe_operation,
            priority=priority,
            session_id=session_id
        )
    
    def execute_sync(self, operation: str, *args, **kwargs):
        """同步执行存储引擎操作"""
        with self._operation_lock:
            method = getattr(self.storage_engine, operation)
            return method(*args, **kwargs)
